#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author: Edward Roualdes
#         University of Kentucky

import sys
from multiprocessing import Pool
from Bio import SeqIO
from sortBarcode import sortBarcode, sortInds
from sortLoci import sortLoci
from sanity import sane
import os.path
import utils
import hap
import structure
import align as a
import phase

def init(argv=None):
    if argv is None:
        argv = sys.argv

    print "\nWelcome to NextAllele\n"
    configfile = utils.parseCommandLine(argv)
    opts = sane(configfile)
    return opts

def main(opts = None):

    if opts['sort']:
        utils.mkdir(
            opts['OUTPUTDIRECTORY'], 
            True,
            opt = ("Please specify a different "
                   "OUTPUTDIRECTORY in the CONFIGFILE "))
        if opts['BARCODEFILE']:
            srt, c = sortBarcode(
                opts['INPUTFILE'], 
                opts['BARCODEFILE'], 
                opts['trim'])
        else:
            srt, c = sortInds(opts['INPUTFILE'])

        srt = sortLoci(srt, opts['BLASTREPORT'], c)
        # write sorted seqs to file
        utils.writeSeqs(srt, opts['OUTPUTDIRECTORY'])
            
    if opts['align']:
        print "Creating primary alignments..."
        p = Pool(processes = opts['cpu'])
        p.map(alignTool, a.files2align(opts['OUTPUTDIRECTORY']))

    if opts['structure'] or opts['phase']:
        # phase haplotypes
        hapDir = hap.callHaps(opts['OUTPUTDIRECTORY'], approx = False) # [2014.05.29] TODO: add option to approx haplotypes
        # align haplotypes
        p = Pool(processes = opts['cpu'])
        p.map(alignTool, a.files2align(hapDir))

    if opts['structure']:
        structure.mkStruct(opts['OUTPUTDIRECTORY'], opts['cleansnp'])

    if opts['phase']:
        phase.mkPhase(opts['OUTPUTDIRECTORY'])
        
    if opts['runphase']:
        global phaseDir 
        phaseDir = os.path.join(opts['OUTPUTDIRECTORY'], 'phase')
        # run program PHASE on all created files
        pp = Pool(processes = opts['cpu'])
        pp.map(phaseTool, phase.files2phase(phaseDir))
            
    print "NextAllele finished."

# appease multiprocessing (and GIL)
def alignTool(x):
    return a.aligner(x, opts["ALIGNTOOL"][0], opts["ALIGNTOOL"][1])

def phaseTool(x):
    return phase.runPhase(x, opts['PHASETOOL'], phaseDir)

if __name__ == '__main__':
    opts = init()
    sys.exit(main(opts))
